from llama_index.core import VectorStoreIndex, Document, Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.openai import OpenAI
from dotenv import load_dotenv
import os

# 1. 加载环境变量
load_dotenv()

# 2. 关键修改：完全绕过模型校验
class PatchedOpenAI(OpenAI):
    @property
    def metadata(self):
        from llama_index.core.llms import LLMMetadata  # 新版路径
        return LLMMetadata(
            context_window=128000,  # GLM-4的上下文长度
            is_chat_model=True,
            model_name="glm-4"
        )

# 3. 配置全局设置
Settings.llm = PatchedOpenAI(
    model="glm-4",
    api_key=os.getenv("ZHIPUAI_API_KEY"),
    api_base="https://open.bigmodel.cn/api/paas/v4",
    is_chat_model=True
)

# 4. 使用本地嵌入模型
Settings.embed_model = HuggingFaceEmbedding(
    model_name="D:/ideaSpace/MyPython/models/bge-small-zh-v1.5"
)

# 5. 测试查询
documents = [Document(text="""《灭神纪∙猢狲》是一款动作角色扮演游戏...
悟空有两种形态变化：
1. 金刚形态：侧重力量型打击
2. 魔佛形态：专注法术攻击""")]

index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine()

response = query_engine.query("游戏中悟空有哪些形态变化？")
print(response)